import torch.nn as nn
from .torch.module_wrapper import NVTXModuleWrapper

class BaseModule():
    supported_submodules = [
        'backbone',
        'neck',
        'head'
    ]
    def __init__(self, module: nn.Module,
                 forward_color: str = 'green',
                 backward_color: str = 'red'):
        self.module = module
        self.forward_color = forward_color
        self.backward_color = backward_color

        self.backbone = None
        self.backbone_hook = None

        self.neck = None
        self.neck_hook = None

        self.head = None
        self.head_hook = None

        self.submodules_conformed = False
        self.submodule_hook_maps = {}
        
    def __call__(self):
        self._annotate()
        return self.module

    def _get_submodules(self):
        raise NotImplementedError

    def _annotate(self):
        self._get_submodules()
        self.submodule_hook_maps = {self.backbone: self.backbone_hook,
                               self.neck: self.neck_hook,
                               self.head: self.head_hook}
        assert self.submodules_conformed, 'submodules not conformed first'
        for submodule, wrapper in self.submodule_hook_maps.items():
            wrapper = NVTXModuleWrapper(submodule)
            wrapper.register()